﻿using System;
using System.Collections.Generic;
using System.Diagnostics;

namespace g4
{
    /// <summary>
    /// 拉普拉斯网格平滑器：通过求解拉普拉斯方程实现网格表面平滑，支持顶点软约束
    /// 核心原理：通过调整顶点位置使网格的拉普拉斯能量最小化，同时可通过约束保留关键特征
    /// </summary>
    public class LaplacianMeshSmoother
    {
        /// <summary>
        /// 待平滑处理的网格对象
        /// </summary>
        public DMesh3 Mesh;

        // 基于网格固定信息的字段（初始化后不变）
        /// <summary>
        /// 压缩存储的稀疏矩阵（拉普拉斯矩阵或其平方，用于线性求解）
        /// </summary>
        PackedSparseMatrix PackedM;
        /// <summary>
        /// 网格顶点数量（活跃顶点，非无效ID）
        /// </summary>
        int N;
        /// <summary>
        /// 索引映射：[线性索引] -> 网格顶点ID（用于将稀疏矩阵索引映射回原始网格顶点）
        /// </summary>
        int[] ToMeshV;
        /// <summary>
        /// 索引映射：[网格顶点ID] -> 线性索引（用于将原始网格顶点映射到稀疏矩阵索引）
        /// </summary>
        int[] ToIndex;
        /// <summary>
        /// 初始顶点坐标的X/Y/Z分量（线性索引对应）
        /// </summary>
        double[] Px, Py, Pz;
        /// <summary>
        /// 每个顶点的邻居数量（用于计算拉普拉斯权重）
        /// </summary>
        int[] nbr_counts;
        /// <summary>
        /// 初始网格位置的拉普拉斯向量（X/Y/Z分量），平滑时会被清零以实现"去噪"效果
        /// </summary>
        double[] MLx, MLy, MLz;

        // 约束相关字段
        /// <summary>
        /// 顶点软约束结构：控制顶点在平滑过程中的位置偏差
        /// </summary>
        public struct SoftConstraintV
        {
            /// <summary>
            /// 约束目标位置（顶点应尽量靠近此位置）
            /// </summary>
            public Vector3d Position;
            /// <summary>
            /// 约束权重（值越大，顶点越难偏离目标位置）
            /// </summary>
            public double Weight;
            /// <summary>
            /// 是否在求解后强制固定到目标位置（硬约束）
            /// </summary>
            public bool PostFix;
        }
        /// <summary>
        /// 软约束字典：key为顶点ID，value为约束信息
        /// </summary>
        Dictionary<int, SoftConstraintV> SoftConstraints = new Dictionary<int, SoftConstraintV>();
        /// <summary>
        /// 是否存在求解后需要强制固定的约束
        /// </summary>
        bool HavePostFixedConstraints = false;


        // 求解相关动态字段（约束更新后需重新计算）
        /// <summary>
        /// 是否需要更新求解相关的矩阵和向量（约束变化时置为true）
        /// </summary>
        bool need_solve_update;
        /// <summary>
        /// 约束权重对角矩阵（仅对角线有值，对应顶点约束权重）
        /// </summary>
        DiagonalMatrix WeightsM;
        /// <summary>
        /// 约束项的右手边向量（X/Y/Z分量），由约束位置和权重计算得到
        /// </summary>
        double[] Cx, Cy, Cz;
        /// <summary>
        /// 线性系统的总右手边向量（拉普拉斯项+约束项，X/Y/Z分量）
        /// </summary>
        double[] Bx, By, Bz;
        /// <summary>
        /// 预处理矩阵（用于加速共轭梯度求解，这里用Jacobi预处理）
        /// </summary>
        DiagonalMatrix Preconditioner;


        /// <summary>
        /// 是否使用软约束的正规方程（A^T A x = A^T b）
        /// 参考：http://sites.fas.harvard.edu/~cs277/papers/deformation_survey.pdf 附录C
        /// </summary>
        public bool UseSoftConstraintNormalEquations = true;


        // 求解结果
        /// <summary>
        /// 平滑后顶点坐标的X/Y/Z分量（线性索引对应）
        /// </summary>
        double[] Sx, Sy, Sz;


        /// <summary>
        /// 构造拉普拉斯网格平滑器
        /// </summary>
        /// <param name="mesh">待平滑的网格（需是紧凑存储的，即无无效顶点ID）</param>
        public LaplacianMeshSmoother(DMesh3 mesh)
        {
            Mesh = mesh;
            // 断言：确保网格是紧凑的（顶点ID连续，无间隙），否则索引映射会出错
            Util.gDevAssert(mesh.IsCompact);
        }


        /// <summary>
        /// 为指定顶点设置软约束
        /// </summary>
        /// <param name="vID">顶点ID</param>
        /// <param name="targetPos">目标位置</param>
        /// <param name="weight">约束权重（越大约束越强）</param>
        /// <param name="bForceToFixedPos">是否在求解后强制固定到目标位置（硬约束）</param>
        public void SetConstraint(int vID, Vector3d targetPos, double weight, bool bForceToFixedPos = false)
        {
            SoftConstraints[vID] = new SoftConstraintV()
            {
                Position = targetPos,
                Weight = weight,
                PostFix = bForceToFixedPos
            };
            HavePostFixedConstraints = HavePostFixedConstraints || bForceToFixedPos;
            need_solve_update = true; // 约束变化，需更新求解参数
        }

        /// <summary>
        /// 检查指定顶点是否有约束
        /// </summary>
        /// <param name="vID">顶点ID</param>
        /// <returns>若有约束则返回true</returns>
        public bool IsConstrained(int vID)
        {
            return SoftConstraints.ContainsKey(vID);
        }

        /// <summary>
        /// 清除所有顶点约束
        /// </summary>
        public void ClearConstraints()
        {
            SoftConstraints.Clear();
            HavePostFixedConstraints = false;
            need_solve_update = true; // 约束变化，需更新求解参数
        }


        /// <summary>
        /// 初始化平滑器：构建拉普拉斯矩阵、索引映射、初始向量等
        /// 这是预处理步骤，需在调用Solve前执行
        /// </summary>
        public void Initialize()
        {
            // 初始化索引映射：将网格顶点ID映射到连续的线性索引（0~N-1）
            ToMeshV = new int[Mesh.MaxVertexID];
            ToIndex = new int[Mesh.MaxVertexID];
            N = 0; // 活跃顶点数量
            foreach (int vid in Mesh.VertexIndices())
            {
                ToMeshV[N] = vid;   // 线性索引N对应原始顶点vid
                ToIndex[vid] = N;   // 原始顶点vid对应线性索引N
                N++;
            }

            // 初始化顶点坐标数组（存储初始位置）
            Px = new double[N];
            Py = new double[N];
            Pz = new double[N];
            nbr_counts = new int[N]; // 存储每个顶点的邻居数量
            SymmetricSparseMatrix M = new SymmetricSparseMatrix(); // 对称稀疏矩阵（拉普拉斯矩阵）

            // 填充初始坐标和邻居数量
            for (int i = 0; i < N; ++i)
            {
                int vid = ToMeshV[i];
                Vector3d v = Mesh.GetVertex(vid);
                Px[i] = v.x;
                Py[i] = v.y;
                Pz[i] = v.z;
                nbr_counts[i] = Mesh.GetVtxEdgeCount(vid); // 边数即邻居数（每个边连接一个邻居）
            }

            // 构建拉普拉斯矩阵（核心步骤）
            for (int i = 0; i < N; ++i)
            {
                int vid = ToMeshV[i];
                int n = nbr_counts[i]; // 当前顶点的邻居数

                double sum_w = 0; // 用于计算对角线元素（负的非对角线元素之和）
                // 遍历当前顶点的所有邻居
                foreach (int nbrvid in Mesh.VtxVerticesItr(vid))
                {
                    int j = ToIndex[nbrvid]; // 邻居顶点的线性索引
                    int n2 = nbr_counts[j];  // 邻居的邻居数

                    // 拉普拉斯权重计算（可选方案）：
                    // 1. 简单均匀权重：double w = -1;
                    // 2. 基于邻居数的调和权重：double w = -1.0 / Math.Sqrt(n + n2);（这里使用此方案）
                    // 3. 基于当前顶点邻居数的均匀权重：double w = -1.0 / n;
                    double w = -1.0 / Math.Sqrt(n + n2);

                    M.Set(i, j, w); // 设置非对角线元素（对称矩阵，i<j时只需设置一次）
                    sum_w += w;     // 累加非对角线元素值
                }
                // 对角线元素 = 负的非对角线元素之和（确保行和为0，拉普拉斯矩阵性质）
                sum_w = -sum_w;
                M.Set(i, i, sum_w);
            }

            // 根据配置构建最终的PackedM矩阵
            if (UseSoftConstraintNormalEquations)
            {
                // 使用拉普拉斯矩阵的平方（M^T * M，因M对称，等价于M*M）
                PackedM = M.SquarePackedParallel();
            }
            else
            {
                // 直接使用拉普拉斯矩阵
                PackedM = new PackedSparseMatrix(M);
            }

            // 计算初始网格位置的拉普拉斯向量（M * P）
            MLx = new double[N];
            MLy = new double[N];
            MLz = new double[N];
            PackedM.Multiply(Px, MLx); // MLx = M * Px
            PackedM.Multiply(Py, MLy); // MLy = M * Py
            PackedM.Multiply(Pz, MLz); // MLz = M * Pz

            // 清零拉普拉斯向量：这是"平滑"的核心——目标是让新位置的拉普拉斯值为0（能量最小）
            for (int i = 0; i < Px.Length; ++i)
            {
                MLx[i] = 0;
                MLy[i] = 0;
                MLz[i] = 0;
            }

            // 初始化求解所需的缓冲和矩阵
            Preconditioner = new DiagonalMatrix(N); // 预处理矩阵（Jacobi）
            WeightsM = new DiagonalMatrix(N);       // 约束权重矩阵
            Cx = new double[N]; Cy = new double[N]; Cz = new double[N]; // 约束项右手边
            Bx = new double[N]; By = new double[N]; Bz = new double[N]; // 总右手边
            Sx = new double[N]; Sy = new double[N]; Sz = new double[N]; // 求解结果

            need_solve_update = true;
            UpdateForSolve(); // 初始化求解参数
        }


        /// <summary>
        /// 更新求解相关的矩阵和向量（约束变化后需调用）
        /// 构建权重矩阵、右手边向量和预处理矩阵
        /// </summary>
        void UpdateForSolve()
        {
            if (need_solve_update == false)
                return;

            // 构建约束权重矩阵和约束项右手边
            WeightsM.Clear();
            Array.Clear(Cx, 0, N);
            Array.Clear(Cy, 0, N);
            Array.Clear(Cz, 0, N);
            foreach (var constraint in SoftConstraints)
            {
                int vid = constraint.Key;
                int i = ToIndex[vid]; // 线性索引
                double w = constraint.Value.Weight;

                // 若使用正规方程，权重需平方（因正规方程中约束项为w^2）
                if (UseSoftConstraintNormalEquations)
                    w = w * w;

                WeightsM.Set(i, i, w); // 权重矩阵对角线元素
                Vector3d pos = constraint.Value.Position;
                // 约束项右手边：C = w * target_pos
                Cx[i] = w * pos.x;
                Cy[i] = w * pos.y;
                Cz[i] = w * pos.z;
            }

            // 总右手边向量 = 拉普拉斯项（已清零） + 约束项
            for (int i = 0; i < N; ++i)
            {
                Bx[i] = MLx[i] + Cx[i];
                By[i] = MLy[i] + Cy[i];
                Bz[i] = MLz[i] + Cz[i];
            }

            // 更新Jacobi预处理矩阵：对角线元素为1/(M_ii + W_ii)
            // 预处理用于加速共轭梯度的收敛速度
            for (int i = 0; i < N; i++)
            {
                double diag_value = PackedM[i, i] + WeightsM[i, i]; // 总矩阵的对角线元素
                Preconditioner.Set(i, i, 1.0 / diag_value);
            }

            need_solve_update = false;
        }


        /// <summary>
        /// 使用多共轭梯度法求解线性系统（X/Y/Z分量独立求解）
        /// </summary>
        /// <param name="Result">输出参数：存储平滑后的顶点坐标（需与网格顶点ID对应）</param>
        /// <returns>若求解成功则返回true</returns>
        public bool SolveMultipleCG(Vector3d[] Result)
        {
            if (WeightsM == null)
                Initialize(); // 若未初始化，强制初始化

            UpdateForSolve(); // 确保求解参数为最新

            // 以初始位置作为求解的初始猜测值
            Array.Copy(Px, Sx, N);
            Array.Copy(Py, Sy, N);
            Array.Copy(Pz, Sz, N);

            // 定义矩阵乘法函数：(M + W) * X，其中M是拉普拉斯矩阵，W是权重矩阵
            Action<double[], double[]> CombinedMultiply = (X, B) => {
                // 计算M * X（并行计算加速）
                PackedM.Multiply_Parallel(X, B);
                // 加上W * X（权重矩阵是对角矩阵，直接相乘）
                for (int i = 0; i < N; ++i)
                    B[i] += WeightsM.D[i] * X[i];
            };

            // 为X/Y/Z分量分别创建共轭梯度求解器
            SparseSymmetricCG SolverX = new SparseSymmetricCG()
            {
                B = Bx,  // 右手边向量
                X = Sx,  // 解向量（初始猜测值）
                MultiplyF = CombinedMultiply,  // 矩阵乘法函数
                PreconditionMultiplyF = Preconditioner.Multiply, // 预处理函数
                UseXAsInitialGuess = true
            };
            SparseSymmetricCG SolverY = new SparseSymmetricCG()
            {
                B = By,
                X = Sy,
                MultiplyF = CombinedMultiply,
                PreconditionMultiplyF = Preconditioner.Multiply,
                UseXAsInitialGuess = true
            };
            SparseSymmetricCG SolverZ = new SparseSymmetricCG()
            {
                B = Bz,
                X = Sz,
                MultiplyF = CombinedMultiply,
                PreconditionMultiplyF = Preconditioner.Multiply,
                UseXAsInitialGuess = true
            };

            // 并行求解三个分量
            SparseSymmetricCG[] solvers = new SparseSymmetricCG[3] { SolverX, SolverY, SolverZ };
            bool[] ok = new bool[3];
            int[] indices = new int[3] { 0, 1, 2 };

            // 并行执行求解（预处理版本收敛更快）
            Action<int> SolveF = (i) => { ok[i] = solvers[i].SolvePreconditioned(); };
            gParallel.ForEach(indices, SolveF);

            // 检查是否所有分量都求解成功
            if (ok[0] == false || ok[1] == false || ok[2] == false)
                return false;

            // 将求解结果映射回原始网格顶点ID
            for (int i = 0; i < N; ++i)
            {
                int vid = ToMeshV[i];
                Result[vid] = new Vector3d(Sx[i], Sy[i], Sz[i]);
            }

            // 应用硬约束：强制固定需要PostFix的顶点
            if (HavePostFixedConstraints)
            {
                foreach (var constraint in SoftConstraints)
                {
                    if (constraint.Value.PostFix)
                    {
                        int vid = constraint.Key;
                        Result[vid] = constraint.Value.Position;
                    }
                }
            }

            return true;
        }


        /// <summary>
        /// 使用多右手边共轭梯度法求解线性系统（X/Y/Z分量联合求解，更高效）
        /// </summary>
        /// <param name="Result">输出参数：存储平滑后的顶点坐标</param>
        /// <returns>若求解成功则返回true</returns>
        public bool SolveMultipleRHS(Vector3d[] Result)
        {
            if (WeightsM == null)
                Initialize(); // 若未初始化，强制初始化

            UpdateForSolve(); // 确保求解参数为最新

            // 初始化多右手边系统：B为3个分量的右手边，X为初始猜测值（初始位置）
            double[][] B = BufferUtil.InitNxM(3, N, new double[][] { Bx, By, Bz });
            double[][] X = BufferUtil.InitNxM(3, N, new double[][] { Px, Py, Pz });

            // 定义多分量矩阵乘法：(M + W) * X，同时处理X/Y/Z
            Action<double[][], double[][]> CombinedMultiply = (Xt, Bt) => {
                // 并行计算M * X的三个分量
                PackedM.Multiply_Parallel_3(Xt, Bt);
                // 加上W * X（并行处理每个分量）
                gParallel.ForEach(Interval1i.Range(3), (j) => {
                    BufferUtil.MultiplyAdd(Bt[j], WeightsM.D, Xt[j]);
                });
            };

            // 定义多分量预处理函数
            Action<double[][], double[][]> CombinedPreconditionerMultiply = (Xt, Bt) => {
                gParallel.ForEach(Interval1i.Range(3), (j) => {
                    Preconditioner.Multiply(Xt[j], Bt[j]);
                });
            };

            // 创建多右手边共轭梯度求解器
            SparseSymmetricCGMultipleRHS Solver = new SparseSymmetricCGMultipleRHS()
            {
                B = B,  // 多右手边向量（3个分量）
                X = X,  // 多解向量（初始猜测值）
                MultiplyF = CombinedMultiply,  // 矩阵乘法函数
                PreconditionMultiplyF = CombinedPreconditionerMultiply, // 预处理函数
                UseXAsInitialGuess = true
            };

            // 执行预处理求解
            bool ok = Solver.SolvePreconditioned();

            if (ok == false)
                return false;

            // 将求解结果映射回原始网格顶点ID
            for (int i = 0; i < N; ++i)
            {
                int vid = ToMeshV[i];
                Result[vid] = new Vector3d(X[0][i], X[1][i], X[2][i]);
            }

            // 应用硬约束
            if (HavePostFixedConstraints)
            {
                foreach (var constraint in SoftConstraints)
                {
                    if (constraint.Value.PostFix)
                    {
                        int vid = constraint.Key;
                        Result[vid] = constraint.Value.Position;
                    }
                }
            }

            return true;
        }


        /// <summary>
        /// 根据网格规模自动选择求解方式（小网格用多CG，大网格用多RHS）
        /// </summary>
        /// <param name="Result">输出参数：存储平滑后的顶点坐标</param>
        /// <returns>若求解成功则返回true</returns>
        public bool Solve(Vector3d[] Result)
        {
            // 顶点数小于10000时，多CG求解更快；反之多RHS更高效
            if (Mesh.VertexCount < 10000)
                return SolveMultipleCG(Result);
            else
                return SolveMultipleRHS(Result);
        }


        /// <summary>
        /// 求解并将平滑结果更新到原始网格
        /// </summary>
        /// <returns>若更新成功则返回true</returns>
        public bool SolveAndUpdateMesh()
        {
            int maxVid = Mesh.MaxVertexID;
            Vector3d[] Result = new Vector3d[maxVid];
            if (Solve(Result) == false)
                return false;
            // 将结果写回网格（只更新有效顶点）
            for (int i = 0; i < maxVid; ++i)
            {
                if (Mesh.IsVertex(i))
                { // 检查顶点是否有效
                    Mesh.SetVertex(i, Result[i]);
                }
            }
            return true;
        }


        /// <summary>
        /// 对网格的指定三角面子集进行局部平滑
        /// 核心逻辑：约束子集边界顶点，对内部顶点进行平滑，同时控制约束强度衰减
        /// </summary>
        /// <param name="mesh">目标网格</param>
        /// <param name="triangles">待平滑的三角面集合</param>
        /// <param name="nConstrainLoops">约束扩散层数（控制边界约束向内部的衰减范围）</param>
        /// <param name="nIncludeExteriorRings">包含的外部环层数（扩展平滑区域）</param>
        /// <param name="bPreserveExteriorRings">是否保留外部环的原始形状</param>
        /// <param name="borderWeight">边界顶点的约束权重</param>
        /// <param name="interiorWeight">内部顶点的基础约束权重（控制整体平滑强度）</param>
        public static void RegionSmooth(DMesh3 mesh, IEnumerable<int> triangles,
            int nConstrainLoops,
            int nIncludeExteriorRings,
            bool bPreserveExteriorRings,
            double borderWeight = 10.0, double interiorWeight = 0.0)
        {
            // 存储需要固定的顶点（外部环顶点）
            HashSet<int> fixedVerts = new HashSet<int>();
            if (nIncludeExteriorRings > 0)
            {
                // 选择并扩展三角面区域
                MeshFaceSelection expandTris = new MeshFaceSelection(mesh);
                expandTris.Select(triangles);
                if (bPreserveExteriorRings)
                {
                    // 保留外部环：标记外部环顶点为固定顶点
                    MeshEdgeSelection bdryEdges = new MeshEdgeSelection(mesh);
                    bdryEdges.SelectBoundaryTriEdges(expandTris); // 选中区域的边界边
                    expandTris.ExpandToOneRingNeighbours(nIncludeExteriorRings); // 扩展区域
                    MeshVertexSelection startVerts = new MeshVertexSelection(mesh);
                    startVerts.SelectTriangleVertices(triangles); // 原始区域顶点
                    startVerts.DeselectEdges(bdryEdges); // 排除边界边顶点
                    MeshVertexSelection expandVerts = new MeshVertexSelection(mesh, expandTris);
                    foreach (int vid in expandVerts)
                    {
                        if (!startVerts.IsSelected(vid)) // 外部环顶点（非原始区域）
                            fixedVerts.Add(vid);
                    }
                }
                else
                {
                    // 不保留外部环：直接扩展区域
                    expandTris.ExpandToOneRingNeighbours(nIncludeExteriorRings);
                }
                triangles = expandTris; // 更新待平滑的三角面
            }

            // 创建区域算子：提取子网格进行局部处理
            RegionOperator region = new RegionOperator(mesh, triangles);
            DSubmesh3 submesh = region.Region; // 子网格（仅包含选中区域）
            DMesh3 smoothMesh = submesh.SubMesh; // 子网格的独立副本
            LaplacianMeshSmoother smoother = new LaplacianMeshSmoother(smoothMesh);

            // 将固定顶点映射到子网格中
            HashSet<int> subFixedVerts = new HashSet<int>();
            foreach (int base_vid in fixedVerts)
                subFixedVerts.Add(submesh.MapVertexToSubmesh(base_vid));

            // 约束子网格的边界顶点（原始区域的边界）
            double w = borderWeight;
            HashSet<int> constrained = (submesh.BaseBorderV.Count > 0) ? new HashSet<int>() : null;
            foreach (int base_vid in submesh.BaseBorderV)
            {
                int sub_vid = submesh.BaseToSubV[base_vid]; // 子网格中的顶点ID
                // 约束到初始位置，权重为borderWeight，且求解后固定
                smoother.SetConstraint(sub_vid, smoothMesh.GetVertex(sub_vid), w, true);
                if (constrained != null)
                    constrained.Add(sub_vid);
            }

            // 扩散约束：向内部扩展约束，权重随层数平方根衰减（控制连续性）
            if (constrained != null && constrained.Count > 0)
            {
                w = Math.Sqrt(w); // 权重衰减
                for (int k = 0; k < nConstrainLoops; ++k)
                {
                    HashSet<int> next_layer = new HashSet<int>(); // 下一层约束顶点

                    foreach (int sub_vid in constrained)
                    {
                        // 遍历当前约束顶点的邻居
                        foreach (int nbr_vid in smoothMesh.VtxVerticesItr(sub_vid))
                        {
                            if (!constrained.Contains(nbr_vid))
                            { // 未被约束的邻居
                                if (!smoother.IsConstrained(nbr_vid))
                                {
                                    // 设置约束，权重为当前w，固定顶点强制固定
                                    smoother.SetConstraint(nbr_vid, smoothMesh.GetVertex(nbr_vid), w, subFixedVerts.Contains(nbr_vid));
                                }
                                next_layer.Add(nbr_vid);
                            }
                        }
                    }

                    constrained.UnionWith(next_layer); // 合并新约束顶点
                    w = Math.Sqrt(w); // 权重继续衰减
                }
            }

            // 为内部顶点添加基础约束（若需要）
            if (interiorWeight > 0)
            {
                foreach (int vid in smoothMesh.VertexIndices())
                {
                    if (!smoother.IsConstrained(vid))
                    {
                        smoother.SetConstraint(vid, smoothMesh.GetVertex(vid), interiorWeight, subFixedVerts.Contains(vid));
                    }
                }
            }
            else if (subFixedVerts.Count > 0)
            {
                // 若有固定顶点但无内部权重，强制固定顶点保持原位
                foreach (int vid in subFixedVerts)
                {
                    if (!smoother.IsConstrained(vid))
                        smoother.SetConstraint(vid, smoothMesh.GetVertex(vid), 0, true);
                }
            }

            // 执行平滑并将结果回传到原始网格
            smoother.SolveAndUpdateMesh();
            region.BackPropropagateVertices(true); // 将子网格的修改同步到原始网格
        }
    }
}